import nltk
import numpy as np
import tensorflow as tf
import random
import _pickle as pickle
from nltk.tokenize import word_tokenize
from nltk.stem import WordNetLemmatizer
from collections import Counter


pos_file = 'pos.txt'
neg_file = 'neg.txt'

# 创建词汇表
def create_lexicon(pos_file, neg_file):
    lex = []

    def process_file(file):
        with open(file, 'r') as f:
            tmps = []
            lines = f.readlines()
            # print(lines)
            for line in lines:
                words = word_tokenize(line.lower())
                tmps += words
            return tmps

    lex += process_file(pos_file)
    lex += process_file(neg_file)
    lemmatizer = WordNetLemmatizer()
    lex = [lemmatizer.lemmatize(word) for word in lex]
    word_count = Counter(lex)
    # print(word_count)
    # {'.': 13944, ',': 10536, 'the': 10120, 'a': 9444, 'and': 7108, 'of': 6624, 'it': 4748, 'to': 3940......}
    # 去掉一些常用词,像the,a and等等，和一些不常用词; 这些词对判断一个评论是正面还是负面没有做任何贡献
    for word in word_count:
        # 去除那些 . # * 还有单个字符的
        # if len(word) == 1:
        #     continue
        if word_count[word] < 2000 and word_count[word] > 20:  # 这写死了，好像能用百分比
            lex.append(word)  # 齐普夫定律-使用Python验证文本的Zipf分布 http://blog.topspeedsnail.com/archives/9546
    return lex


# 把每条评论转换为向量, 转换原理：
# 假设lex为['woman', 'great', 'feel', 'actually', 'looking', 'latest', 'seen', 'is'] 当然实际上要大的多
# 评论'i think this movie is great' 转换为 [0,1,0,0,0,0,0,1], 把评论中出现的字在lex中标记，出现过的标记为1，其余标记为0
def normailize_datasets(lex):
    # lex:词汇表；review:评论；clf:评论对应的分类，[0,1]代表负面评论 [1,0]代表正面评论
    dataset = []

    # lex:词汇表；review:评论；clf:评论对应的分类，[0,1]代表负面评论 [1,0]代表正面评论
    def string_to_vector(lex, review, clf):
        words = word_tokenize(review.lower())
        lemmatizer = WordNetLemmatizer()
        words = [lemmatizer.lemmatize(word) for word in words]

        features = np.zeros(len(lex))
        # 一个句子中某个词可能出现两次,可以用+=1，其实区别不大
        # [features[lex.index(word)] for word in words if word in lex]
        for word in words:
            if word in lex:
                features[lex.index(word)] = 1
        return [features, clf]

    with open(pos_file, 'r') as f:
        lines = f.readlines()
        # next(f)
        for line in lines:
            one_sample = string_to_vector(lex, line, [1, 0])  # [array([ 0.,  1.,  0., ...,  0.,  0.,  0.]), [1,0]]
            dataset.append(one_sample)
    with open(neg_file, 'r') as f:
        lines = f.readlines()
        # next(f)
        for line in lines:
            one_sample = string_to_vector(lex, line, [0, 1])  # [array([ 0.,  0.,  0., ...,  0.,  0.,  0.]), [0,1]]]
            dataset.append(one_sample)

    # print(len(dataset))
    return dataset


lex = create_lexicon(pos_file, neg_file)
dataset = normailize_datasets(lex)
random.shuffle(dataset)
#把整理好的数据保存到文件，方便使用。到此完成了数据的整理工作
# with open('save.pickle','wb') as f:
#     pickle.dump(dataset, f)
# 取样本中的10%做为测试数据
test_size = int(len(dataset) * 0.1)

# 转换成na_array
dataset = np.array(dataset)

# 定义训练集
train_set = dataset[test_size:]

# 定义测试集
test_set = dataset[0:test_size]


# 定义每层的神经元 (输入 输出 隐藏)
n_input_layer = len(lex)

n_hide_layer1 = 1000
n_hide_layer2 = 1000

n_output_layer = 2

# 定义待训练的神经网络
def neural_network(data):
    # 定义第一层"神经元"的权重和biases
    layer_1_w_b = {'w_': tf.Variable(tf.random_normal([n_input_layer, n_hide_layer1])),
                   'b_': tf.Variable(tf.random_normal([n_hide_layer1]))}
    # 定义第二层"神经元"的权重和biases
    layer_2_w_b = {'w_': tf.Variable(tf.random_normal([n_hide_layer1, n_hide_layer2])),
                   'b_': tf.Variable(tf.random_normal([n_hide_layer2]))}
    # 定义输出层"神经元"的权重和biases
    layer_output_w_b = {'w_': tf.Variable(tf.random_normal([n_hide_layer2, n_output_layer])),
                        'b_': tf.Variable(tf.random_normal([n_output_layer]))}

    # w·x+b
    layer_1 = tf.add(tf.matmul(data, layer_1_w_b['w_']), layer_1_w_b['b_'])
    layer_1 = tf.nn.relu(layer_1)  # 激活函数
    layer_2 = tf.add(tf.matmul(layer_1, layer_2_w_b['w_']), layer_2_w_b['b_'])
    layer_2 = tf.nn.relu(layer_2)  # 激活函数
    layer_output = tf.add(tf.matmul(layer_2, layer_output_w_b['w_']), layer_output_w_b['b_'])

    return layer_output

# 每次使用的训练数据量
batch_size = 50

X = tf.placeholder('float',[None, len(train_set[0][0])])
Y = tf.placeholder('float')

# 开始训练
def train_neural_network(X, Y):
    prediction = neural_network(X)
    loss_fun = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(labels=Y, logits=prediction))
    optimizer = tf.train.AdamOptimizer().minimize(loss_fun)

    epochs = 13
    with tf.Session() as session:
        session.run(tf.initialize_all_variables())


        random.shuffle(train_set)
        train_x = dataset[:, 0]
        train_y = dataset[:, 1]
        for epoch in range(epochs):
            i = 0
            epoch_loss = 0
            while i < len(train_x):
                start = i
                end = i + batch_size

                batch_x = train_x[start:end]
                batch_y = train_y[start:end]

                _, c = session.run([optimizer, loss_fun], feed_dict={X: list(batch_x), Y: list(batch_y)})
                epoch_loss += c
                i += batch_size

            print(epoch, ' : ', epoch_loss)

        text_x = test_set[:, 0]
        text_y = test_set[:, 1]
        correct = tf.equal(tf.argmax(prediction, 1), tf.argmax(Y, 1))
        accuracy = tf.reduce_mean(tf.cast(correct, 'float'))
        print('准确率: ', accuracy.eval({X: list(text_x), Y: list(text_y)}))

train_neural_network(X, Y)







